import numpy as np

class Layer:
    def __init__(self):
        self.params = []
        self.grads = []
        self.x = None
        self.cache = None
        
    def forward(self, x):
        raise NotImplementedError
    
    def backward(self, dout):
        raise NotImplementedError
    
class Sigmoid(Layer):
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        out = 1 / (1 + np.exp(-x))
        self.cache = out
        return out
    
    def backward(self, dout):
        cache = self.cache
        dx = dout * cache * (1 - cache)
        return dx
    
class ReLU(Layer):
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        self.mask = x <= 0
        out = x.copy()
        out[self.mask] = 0  
        return out
    
    def backward(self, dout):
        dout[self.mask] = 0
        dx = dout
        return dx

class Affine(Layer):
    def __init__(self, W, b):
        super().__init__()
        self.W = W
        self.b = b
        self.x = None
        self.original_x_shape = None
        out_dim, in_dim = W.shape
        
        self.params = [W, b]
        self.grads = [np.zeros_like(W), np.zeros_like(b)]
        
    def forward(self, x):
        self.original_x_shape = x.shape
        x = x.reshape(x.shape[0], -1)
        self.x = x
        out = np.dot(self.x, self.W.T) + self.b
        return out  
    
    def backward(self, dout):
        dx = np.dot(dout, self.W)
        self.grads[0] = np.dot(self.x.T, dout)
        self.grads[1] = np.sum(dout, axis=0)
        return dx

def softmax(x):
    if x.ndim == 2:
        x = x.T
        x = x - np.max(x, axis=0)
        y = np.exp(x) / np.sum(np.exp(x), axis=0)
        return y.T

def cross_entropy_error(y, t):
    if y.ndim == 1:
        t = t.reshape(1, t.size)
        y = y.reshape(1, y.size)
        
    if t.size == y.size:
        t = t.argmax(axis=1)    
    batch_size = y.shape[0]
    return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size
    
class SoftmaxWithLoss(Layer):
    def __init__(self, loss_func=cross_entropy_error):
        super().__init__()
        self.loss = None
        self.y = None
        self.t = None
        self.loss_func = loss_func
        
    def forward(self, x, t):
        self.t = t
        self.y = softmax(x)
        self.loss = self.loss_func(self.y, self.t)
        return self.y
    
    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        dx = (self.y - self.t) / batch_size
        return dx
